// Copyright 2011 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "google/protobuf-expectation/proto-kit-generator.h"

#include <string.h>
#include <set>
#include <string>

#include "google/protobuf-expectation/stubs/casts.h"
#include "gflags/gflags.h"
#include "glog/logging.h"
#include "google/protobuf-expectation/stubs/scoped_ptr.h"
#include "google/protobuf-expectation/descriptor-proto-printer.h"
#include "google/protobuf/io/zero_copy_stream.h"
#include "google/protobuf/descriptor.pb.h"
#include "google/protobuf/descriptor.h"

DEFINE_int32(at_index_tag_number, 999,
             "The tag number of at_index fields in proto kit.");

using std::set;
using std::string;

using google::protobuf::Descriptor;
using google::protobuf::DescriptorProto;
using google::protobuf::FieldDescriptor;
using google::protobuf::FieldDescriptorProto;
using google::protobuf::FileDescriptor;
using google::protobuf::FileDescriptorProto;
using google::protobuf::Message;
using google::protobuf::compiler::GeneratorContext;
using google::protobuf::io::ZeroCopyOutputStream;

namespace google {
namespace protobuf_expectation {

namespace {

const char kProtoKitFileName[] =
"google/protobuf-expectation/proto-kit.proto";
const char kValueKitTypeName[] =
    ".google.protobuf_expectation.ValueKit";
const char kRepeatedValueKitTypeName[] =
    ".google.protobuf_expectation.RepeatedValueKit";

string AddSuffixToBasename(const string& filename, const string& suffix) {
  string::size_type slash = filename.rfind('/');
  string::size_type dot = filename.rfind('.');
  if ((slash == string::npos && dot != string::npos) ||
      (slash != string::npos && dot != string::npos && slash < dot)) {
    string s(filename);
    return s.insert(dot, suffix);
  } else {
    return filename + suffix;
  }
}

// Returns a proto kit file name for the proto file name.
inline string ToProtoKitFileName(const string& name) {
  return AddSuffixToBasename(name, "-kit");
}

// Returns a proto kit type name for the message type name (without package
// name).
inline string ToProtoKitTypeName(const string& name) {
  return name + "Kit";
}

// Returns a repeated proto kit type name for the message type name (without
// package name).
inline string ToRepeatedProtoKitTypeName(const string& name) {
  return "Repeated" + name + "Kit";
}

// Gets a full type name of a proto kit for the message type.
string GetProtoKitTypeName(const Descriptor* descriptor) {
  string type_name = ToProtoKitTypeName(descriptor->name());
  for (const Descriptor* d = descriptor->containing_type();
       d != NULL; d = d->containing_type()) {
    type_name = ToProtoKitTypeName(d->name()) + "." + type_name;
  }
  return "." + descriptor->file()->package() + "." + type_name;
}

// Gets a full type name of a repeated proto kit for the message type.
string GetRepeatedProtoKitTypeName(const Descriptor* descriptor) {
  string type_name = ToRepeatedProtoKitTypeName(descriptor->name());
  for (const Descriptor* d = descriptor->containing_type();
       d != NULL; d = d->containing_type()) {
    type_name = ToProtoKitTypeName(d->name()) + "." + type_name;
  }
  return "." + descriptor->file()->package() + "." + type_name;
}

// Adds a descriptor proto to the message's field which is specified by the
// field_name_to_add.
DescriptorProto* AddDescriptorProto(
    const string& field_name_to_add, Message* message) {
  const FieldDescriptor* field =
      message->GetDescriptor()->FindFieldByName(field_name_to_add);
  CHECK_NOTNULL(field);
  CHECK(field->is_repeated());
  CHECK_EQ(field->type(), FieldDescriptor::TYPE_MESSAGE);
  CHECK_EQ(field->message_type(), DescriptorProto::descriptor());
  return google::protobuf_expectation::down_cast<DescriptorProto*>(
      message->GetReflection()->AddMessage(message, field));
}

void WriteStringToOutputStream(const string& str,
                               ZeroCopyOutputStream* output) {
  string::size_type pos = 0;
  void* buffer;
  int size;
  while (output->Next(&buffer, &size)) {
    const int bytes = str.size() - pos;
    if (bytes > size) {
      memcpy(buffer, str.data() + pos, size);
      pos += size;
    } else {
      memcpy(buffer, str.data() + pos, bytes);
      output->BackUp(size - bytes);
      break;
    }
  }
}

}  // namespace

bool ProtoKitGenerator::Generate(const FileDescriptor* file,
                                 const string& unused_parameter,
                                 GeneratorContext* generator_context,
                                 string* error) const {
  CHECK_NOTNULL(generator_context);
  CHECK_NOTNULL(error);
  FileDescriptorProto file_descriptor_proto;
  const bool result = ProcessFileDescriptor(file, &file_descriptor_proto,
                                            error);
  if (result) {
    string file_content;
    if (!PrintFileDescriptorProtoToString(file_descriptor_proto,
                                          "proto-kit-generator",
                                          &file_content)) {
      return false;
    }
    scoped_ptr<ZeroCopyOutputStream> output(
        generator_context->Open(ToProtoKitFileName(file->name())));
    WriteStringToOutputStream(file_content, output.get());
  }
  return result;
}

bool ProtoKitGenerator::GenerateDescriptorProto(const FileDescriptor* file,
                                                FileDescriptorProto* output,
                                                string* error) const {
  CHECK_NOTNULL(output);
  CHECK_NOTNULL(error);
  return ProcessFileDescriptor(file, output, error);
}

// Processes the file_descriptor and generates a file descriptor proto of
// proto kit.
bool ProtoKitGenerator::ProcessFileDescriptor(
    const FileDescriptor* file_descriptor,
    FileDescriptorProto* output,
    string* error) const {
  output->set_name(ToProtoKitFileName(file_descriptor->name()));
  output->set_package(file_descriptor->package());
  set<string> deps;
  for (int i = 0; i < file_descriptor->dependency_count(); ++i) {
    deps.insert(ToProtoKitFileName(file_descriptor->dependency(i)->name()));
  }
  deps.insert(kProtoKitFileName);
  for (set<string>::const_iterator it = deps.begin(); it != deps.end(); ++it) {
    output->add_dependency(*it);
  }
  for (int i = 0; i < file_descriptor->message_type_count(); ++i) {
    if (!ProcessDescriptor(file_descriptor->message_type(i), "message_type",
                           output, error)) {
      return false;
    }
  }
  return true;
}

// Processes the descriptor and add one proto kit to describe expectations of
// the singluar message and another proto kit to describe expectations of
// the repeated message.
// The output should be either FileDescriptorProto message with
// field_name_to_add = "message_type" or DescriptorProto message with
// field_name_to_add = "nested_type".
bool ProtoKitGenerator::ProcessDescriptor(
    const Descriptor* descriptor,
    const string& field_name_to_add,
    Message* output,
    string* error) const {
  DescriptorProto* proto_kit =
      AddDescriptorProto(field_name_to_add, output);
  if (!FillProtoKit(descriptor, proto_kit, error)) {
    return false;
  }
  DescriptorProto* repeated_proto_kit =
      AddDescriptorProto(field_name_to_add, output);
  if (!FillRepeatedProtoKit(descriptor, repeated_proto_kit, error)) {
    return false;
  }
  for (int i = 0; i < descriptor->nested_type_count(); ++i) {
    if (!ProcessDescriptor(descriptor->nested_type(i), "nested_type",
                           proto_kit, error)) {
      return false;
    }
  }
  return true;
}

bool ProtoKitGenerator::FillProtoKit(
    const Descriptor* descriptor,
    DescriptorProto* output,
    string* error) const {
  output->set_name(ToProtoKitTypeName(descriptor->name()));
  for (int i = 0; i < descriptor->field_count(); ++i) {
    if (!ProcessFieldDescriptor(descriptor->field(i),
                                output->add_field(), error)) {
      return false;
    }
  }
  FieldDescriptorProto* at_index = output->add_field();
  at_index->set_name("at_index");
  at_index->set_number(FLAGS_at_index_tag_number);
  at_index->set_label(FieldDescriptorProto::LABEL_OPTIONAL);
  at_index->set_type(FieldDescriptorProto::TYPE_MESSAGE);
  at_index->set_type_name(kValueKitTypeName);
  return true;
}

bool ProtoKitGenerator::FillRepeatedProtoKit(
    const Descriptor* descriptor,
    DescriptorProto* output,
    string* error) const {
  output->set_name(ToRepeatedProtoKitTypeName(descriptor->name()));
  FieldDescriptorProto* field_size = output->add_field();
  field_size->set_name("field_size");
  field_size->set_number(1);
  field_size->set_label(FieldDescriptorProto::LABEL_OPTIONAL);
  field_size->set_type(FieldDescriptorProto::TYPE_MESSAGE);
  field_size->set_type_name(kValueKitTypeName);
  FieldDescriptorProto* for_each = output->add_field();
  for_each->set_name("for_each");
  for_each->set_number(2);
  for_each->set_label(FieldDescriptorProto::LABEL_OPTIONAL);
  for_each->set_type(FieldDescriptorProto::TYPE_MESSAGE);
  for_each->set_type_name(ToProtoKitTypeName(descriptor->name()));
  FieldDescriptorProto* have = output->add_field();
  have->set_name("have");
  have->set_number(3);
  have->set_label(FieldDescriptorProto::LABEL_REPEATED);
  have->set_type(FieldDescriptorProto::TYPE_MESSAGE);
  have->set_type_name(ToProtoKitTypeName(descriptor->name()));
  FieldDescriptorProto* do_not_have = output->add_field();
  do_not_have->set_name("do_not_have");
  do_not_have->set_number(4);
  do_not_have->set_label(FieldDescriptorProto::LABEL_REPEATED);
  do_not_have->set_type(FieldDescriptorProto::TYPE_MESSAGE);
  do_not_have->set_type_name(ToProtoKitTypeName(descriptor->name()));
  return true;
}

bool ProtoKitGenerator::ProcessFieldDescriptor(
    const FieldDescriptor* field_descriptor,
    FieldDescriptorProto* output,
    string* error) const {
  CHECK_NE(field_descriptor->number(), FLAGS_at_index_tag_number)
      << "There's a conflict in a field tag number and at_index tag number.";
  const bool is_repeated = field_descriptor->is_repeated();
  const bool is_message =
      field_descriptor->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE;
  if (!is_repeated) {
    if (!is_message) {
      output->set_name(field_descriptor->name());
      output->set_number(field_descriptor->number());
      output->set_label(FieldDescriptorProto::LABEL_OPTIONAL);
      output->set_type(FieldDescriptorProto::TYPE_MESSAGE);
      output->set_type_name(kValueKitTypeName);
    } else {
      output->set_name(field_descriptor->name());
      output->set_number(field_descriptor->number());
      output->set_label(FieldDescriptorProto::LABEL_OPTIONAL);
      output->set_type(FieldDescriptorProto::TYPE_MESSAGE);
      output->set_type_name(
          GetProtoKitTypeName(field_descriptor->message_type()));
    }
  } else {
    if (!is_message) {
      output->set_name(field_descriptor->name());
      output->set_number(field_descriptor->number());
      output->set_label(FieldDescriptorProto::LABEL_OPTIONAL);
      output->set_type(FieldDescriptorProto::TYPE_MESSAGE);
      output->set_type_name(kRepeatedValueKitTypeName);
    } else {
      output->set_name(field_descriptor->name());
      output->set_number(field_descriptor->number());
      output->set_label(FieldDescriptorProto::LABEL_OPTIONAL);
      output->set_type(FieldDescriptorProto::TYPE_MESSAGE);
      output->set_type_name(
          GetRepeatedProtoKitTypeName(field_descriptor->message_type()));
    }
  }
  return true;
}

}  // namespace protobuf_expectation
}  // namespace google
